import numpy as np
from matplotlib import pyplot as plt
from sklearn.cluster import KMeans


def get_cluster(n_clusters=8, max_iter=300):
    km = KMeans(n_clusters=n_clusters, max_iter=max_iter)
    return km


def cluster_plot_2d(X, km_fit):
    plt.figure(figsize=(8, 6))
    plt.subplot(1, 1, 1)
    label_pred = km_fit.labels_  # 获取聚类标签
    centroids = km_fit.cluster_centers_  # 获取聚类中心
    inertia = km_fit.inertia_  # 获取聚类准则的总和
    mark = ['or', 'ob', 'og', 'ok', '^r', '+r', 'sr', 'dr', '<r', 'pr']  # 这里'or'代表中的'o'代表画圈，'r'代表颜色为红色，后面的依次类推
    for i in range(len(label_pred)):
        plt.plot([X[i, 0]], [X[i, 1]], mark[label_pred[i]], markersize=5)
    for centroid in centroids:
        plt.scatter(centroid[0], centroid[1], c='black', marker='x')
    # plt.legend()
    plt.show()


X = np.array([
    [1, 1],
    [1, 1.1],
    [1.1, 1],
    [1, 1.2],
    [1.2, 2],
    [4, 4],
    [4, 4.2],
    [4.2, 4],
    [5, 4]
])
km = KMeans(n_clusters=2, max_iter=200)
km_fit = km.fit(X)
P = km_fit.predict(X)
print(P)

cluster_plot_2d(X, km_fit)
